import os
import pickle
import torch
from deepctr_torch.models import WDL

with open("train_cached_minmax_分桶-1000.pkl", "rb")as f1:
    train_cached = pickle.load(f1)
with open("dev_test_cached_minmax_分桶-1000.pkl", "rb")as f1:
    dev_test_cached = pickle.load(f1)

DEVICE = "cuda:1" if torch.cuda.is_available() else 'cpu'

top10 = [5, 4, 1, 8, 22, 20, 29, 33, 15, 25]
label_weight_norm_manul = [0] + [0.55/10 if i in top10 else 0.45/(121-10) for i in range(1,122)]
label_weight_norm_manul = [100*i for i in label_weight_norm_manul]

model_path = "./save_WDL_minmax-分桶-1000-focal-weight64/"
if not os.path.exists(model_path): os.mkdir(model_path)
model_path += "WDL_prelu-512-256.pt"
model = WDL(linear_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_hidden_units=(512, 256), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path)
model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'], alpha=label_weight_norm_manul)

model.fit(x=train_cached["train_model_input"], y=train_cached["df_train_target_values"],
        validation_data = (dev_test_cached["dev_model_input"], dev_test_cached["df_dev_target_values"]),
        batch_size=512, epochs=40, verbose=1)

